package com.example.dynamicdatabasedemo.config;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Description:
 * @Author: dengping.xu
 * @Data: 2021-07-23
 */
public class DynamicDataSourceHolder {
    private static final Map<String, String> DATABASES = new ConcurrentHashMap<>();
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();

    public static String getDataSourceKey() {
        return contextHolder.get();
    }

    static void clearDataSourceKey() {
        contextHolder.remove();
    }

    public static void switchDb() {
        setDataSourceKey("default");
    }

    public static void switchDb(long storageId) {
        setDataSourceKey(String.valueOf(storageId));
    }

    public static void switchDb(String database) {
        DATABASES.forEach((key, value) -> {
            if (database.equals(value)) {
                setDataSourceKey(key);
            }
        });
    }

    static void setDataSourceKey(String key) {
        if (!DynamicDataSource.isExistDataSource(key)) {
            Map<Long, DbMoudleConfig> errorDbMap = DynamicDataSource.getErrorDataSourceMap();
            Long id = Long.valueOf(key);
            if(!errorDbMap.containsKey(id)){
                throw new IllegalArgumentException("The database is not initialized, storageId = " + id);
            }

            DbMoudleConfig dbMoudleConfig = errorDbMap.get(id);
            DynamicDataSource dynamicDataSource = DynamicDataSource.getInstance();
            DataSource dataSource = DynamicDataSource.createDataSource(dbMoudleConfig);
            if (dataSource == null){
                throw new RuntimeException("The database is disabled");
            }

            Map<Object, Object> dataSourceMap = dynamicDataSource.getDataSourceMap();
            dataSourceMap.put(key,dataSource);
            addSchema(key,dbMoudleConfig.getDatabase());
            dynamicDataSource.setTargetDataSources(dataSourceMap);
            dynamicDataSource.afterPropertiesSet();
        }

        contextHolder.set(key);
    }

    static void addSchema(String key, String database) {
        if(!DATABASES.containsKey(key)){
            DATABASES.put(key, database);
        }
    }
}
